"use client";

import {
  AbsoluteCenter,
  Accordion,
  AccordionButton,
  AccordionIcon,
  AccordionItem,
  AccordionPanel,
  Badge,
  Box,
  Button,
  Card,
  CardBody,
  CircularProgress,
  Divider,
  FormControl,
  Heading,
  Icon,
  IconButton,
  Input,
  Text,
} from "@chakra-ui/react";
import { json } from "@codemirror/lang-json";
import Form from "@rjsf/chakra-ui";
import validator from "@rjsf/validator-ajv8";
import CodeMirror from "@uiw/react-codemirror";
import Ajv from "ajv";
import addFormats from "ajv-formats";
import { useRouter } from "next/navigation";
import React from "react";

import { ChatBubbleBottomCenterTextIcon } from "@heroicons/react/24/outline";
import { useMutation } from "@tanstack/react-query";
import { suggestExtractor, useCreateExtractor } from "../utils/api";

const ajv = new Ajv();
// Adds support for parsing format types like "date-time"
// and "email" in JSON Schema.
// A lot of the JSON Schema generated by LLMS will
// be generated with these formats out of the box.
addFormats(ajv);

/**
 * Component to create a new extractor with fields
 * for name, description, schema, and examples
 */
const CreateExtractor = ({}) => {
  const startSchema = "{}";
  // You might use a mutation hook here if you're
  // using something like React Query for state management
  const [schema, setSchema] = React.useState(startSchema);
  const [creatable, setCreatable] = React.useState(false);
  const [lastValidSchema, setLastValidSchema] = React.useState(
    JSON.parse(startSchema),
  );
  const [currentSchemaValid, setCurrentSchemaValid] = React.useState(true);
  const [userInput, setUserInput] = React.useState("");

  const suggestMutation = useMutation({
    mutationFn: suggestExtractor,
    onSuccess: (data) => {
      let prettySchema = data.json_schema;

      try {
        prettySchema = JSON.stringify(JSON.parse(data.json_schema), null, 2);
      } catch (e) {}

      setSchema(prettySchema);
    },
  });

  const { push } = useRouter();
  const { mutate } = useCreateExtractor({
    onSuccess: (data) => {
      push(`/e/${data.uuid}`);
    },
  });

  React.useMemo(() => {
    try {
      const parsedSchema = JSON.parse(schema);
      ajv.compile(parsedSchema);
      setCurrentSchemaValid(true);
      setLastValidSchema(parsedSchema);
      // OK to create if schema is parseable and not empty
      // and contains an object at the top level
      setCreatable(parsedSchema.type === "object");
    } catch (e) {
      setCurrentSchemaValid(false);
      setCreatable(false);
    }
  }, [schema]);

  const handleSubmit = (event: React.FormEvent<HTMLFormElement>) => {
    event.preventDefault();
    const instruction = "";
    const objectSchema = JSON.parse(schema);
    // Extract information from schema like name, and description
    const name = objectSchema.title || "Unnamed";
    const description = objectSchema.description || "";
    // backend uses varchar(100) for description
    const shortDescription =
      description.length > 100
        ? description.substring(0, 95) + "..."
        : description;

    mutate({
      name,
      description: shortDescription,
      schema: objectSchema,
      instruction,
    });
  };

  const handleSuggest = (event: React.FormEvent<HTMLFormElement>) => {
    event.preventDefault();
    const description = event.currentTarget.userInput.value;
    if (description === "") {
      return;
    }
    suggestMutation.mutate({ description, jsonSchema: schema });
    setUserInput("");
  };

  return (
    <div className="w-4/5 m-auto">
      <Heading size={"md"} className="m-auto w-4/5" textAlign={"center"}>
        What would you like to extract today?
      </Heading>
      <form className="m-auto flex gap-2 mt-5" onSubmit={handleSuggest}>
        <FormControl id="userInput">
          <Input
            htmlSize={4}
            width="100%"
            autoFocus
            height="auto"
            placeholder="Describe your extraction task..."
            value={userInput}
            onChange={(event) => setUserInput(event.target.value)}
          />
        </FormControl>
        {suggestMutation.isPending ? (
          <CircularProgress isIndeterminate />
        ) : (
          <IconButton
            type="submit"
            icon={<Icon as={ChatBubbleBottomCenterTextIcon} />}
            aria-label="OK"
            colorScheme={userInput === "" ? "gray" : "blue"}
            disabled={userInput === ""}
          />
        )}
      </form>
      <form
        className="m-auto flex flex-col content-between gap-5 mt-10"
        onSubmit={handleSubmit}
      >
        <Box position="relative" padding="10">
          <Divider />
          <AbsoluteCenter bg="white" px="10">
            OR
          </AbsoluteCenter>
        </Box>
        <Accordion allowToggle={true}>
          <AccordionItem>
            <AccordionButton>
              Edit JSON Schema
              <div className="ml-auto">
                {currentSchemaValid ? (
                  <Badge colorScheme="green">OK</Badge>
                ) : (
                  <Badge colorScheme="red">Errors!</Badge>
                )}
                <AccordionIcon />
              </div>
            </AccordionButton>
            <AccordionPanel>
              <FormControl isInvalid={!currentSchemaValid}>
                <CodeMirror
                  id="schema"
                  value={schema}
                  aria-label="JSON Schema Editor"
                  onChange={(value) => setSchema(value)}
                  basicSetup={{ autocompletion: true }}
                  extensions={[json()]}
                  minHeight="300px"
                  className="border-4 border-slate-300 border-double"
                />
              </FormControl>
            </AccordionPanel>
          </AccordionItem>
        </Accordion>
        {Object.keys(lastValidSchema).length !== 0 && (
          <>
            <Heading size="md">Preview</Heading>
            {!currentSchemaValid && (
              <Text color="red.500">
                JSON Schema has errors. Showing previous valid JSON Schema.
              </Text>
            )}
            <Card>
              <CardBody>
                <Form
                  schema={lastValidSchema}
                  validator={validator}
                  disabled={!currentSchemaValid || suggestMutation.isPending}
                >
                  {true} {/* Disables the submit button */}
                </Form>
              </CardBody>
            </Card>
          </>
        )}
        <Button className="btn" type="submit" size="lg" isDisabled={!creatable}>
          Create
        </Button>
      </form>
    </div>
  );
};

export default CreateExtractor;
